import numpy as np
import scipy as sp
import help_func
import math

# LevenMarq
class BraynAlg:
    def __init__(self, simData,mxdata):
        self.bd = help_func.BraynData        
        self.kernel = np.mat(simData.kp)
        # SVD
        [V, s, U] = np.linalg.svd(self.kernel, full_matrices=False)
        print "SVD Diff -", np.linalg.norm(self.kernel - np.dot(V, np.dot(np.diag(s), U)), 'fro')
        # chop small sungular values
        lsvd = np.max(s)
        eps = lsvd * mxdata.svdCutOff
        s = s[np.where(s > eps)]
        numSV = np.size(s)
        print "Num Of singular Values - ", numSV 
        self.bd.V = np.mat(V[:, 0:numSV])
        self.bd.U = np.mat(U[0:numSV, :].T)
        self.bd.S = np.mat(np.diag(s))
        print "Kernel Chop Diff- ", np.linalg.norm(self.kernel - self.bd.V * self.bd.S * self.bd.U.T, 'fro')     
        self.bd.data = np.mat(simData.datap).T
        #self.bd.V * self.bd.S * self.bd.U.T    
        self.normK = self.kernel.T * self.kernel        
        self.md = np.ones(self.bd.U.T.shape[1])
        #self.md[0:0] = 1E-100
        self.md /= np.sum(self.md)
        self.md = np.mat(self.md).T
        self.convStep = 100
        self.convMin = 1E-7
        self.maxIter = 5000
        self.u = np.mat(np.ones(self.bd.S.shape[0])).T
        self.ndat = float(self.bd.data.size);
    def updateAlpha(self, alpha):
        self.alpha = alpha
    def calcf(self):
        expVal = np.exp(self.bd.U * self.u)
        self.f = np.multiply(self.md , expVal)        
    def calcdu(self):
        factor = 1
        # Calc f
        self.calcf()
        # Calc K
        diagf = np.mat(np.diag(np.squeeze(np.asarray(self.f))))
        self.K = self.bd.U.T * diagf * self.bd.U
        # Calc A 
        A = (self.mu + self.alpha) * np.eye(self.u.shape[0]) + factor * self.bd.S * self.bd.S * self.K
        #calc g
        preFactor = self.bd.S * self.bd.V.T
        innerP = self.kernel * self.f 
        innerO = self.bd.data
        self.g = factor * preFactor * (innerP - innerO)
        # Calc b
        b = -self.alpha * self.u - self.g
        du = np.mat(np.linalg.lstsq(A, b)[0])	
        return du
    def calcStepSize(self, du):
        dstep = du.T * self.K * du
        return dstep
    def calcEntropy(self):        
        f = np.array(self.f)
        eps = 1E-50
        f[f < eps] = eps
        md = np.array(self.md)
        return np.sum(f - md - f * np.log(f / md))
    def calcchisqr(self):
        chisqr =  np.sum(np.power(self.kernel * self.f - self.bd.data, 2))  
        return chisqr     
    def calcConv(self):     
        gradS = -self.alpha * self.K * self.u  	
        gradL = self.K * self.g	
        diff = gradS - gradL	
        diffabs = math.sqrt(diff.T * diff)
        gradSabs = math.sqrt(gradS.T * gradS)
        gradLabs = math.sqrt(gradL.T * gradL)
        sumg = math.sqrt(gradSabs ** 2 + gradLabs ** 2)
        return diffabs / sumg     
    def sumResults(self):
        bres = help_func.BraynResult() 
        #self.alpha *= 2. / self.ndat        
        self.calcf()
        bres.f = self.f
        #Calc Q
        bres.entropy = self.calcEntropy()   
        bres.chisqr = self.calcchisqr()
        print "chisqr -", bres.chisqr, "alpha S - ", self.alpha * bres.entropy 
        bres.Q = (self.alpha * bres.entropy - bres.chisqr) 
        ## prob alpha
        sqrtf = np.sqrt(self.f)
        sqrtf = np.mat(np.diagflat(sqrtf))
        numFreq = self.f.size
        lambdaMat = sqrtf * self.normK * sqrtf + self.alpha * np.eye(numFreq)     
        L = np.linalg.cholesky(lambdaMat)
        ld = np.sum(np.log(np.power(np.diag(L, 0), 2)))                        
        # print np.flipud(lambdaVal)[0:20]
        bres.logPart = 0.5 * (numFreq * math.log(self.alpha) - ld)
        #bres.ng = np.sum((lambdaVal / (self.alpha + lambdaVal)))
        #print "2 alpha entropy = ", 2 * self.alpha * bres.entropy , "ng ", bres.ng , "entropy = ", bres.entropy
        #print "diff = ", 2 * self.alpha * bres.entropy + bres.ng
        bres.probAlpha = bres.logPart + bres.Q
        ## alpha
        bres.alpha = self.alpha 
        ## calc nabQ
        eps = 1E-50
        bres.f[bres.f < eps] = eps
        nabTemp=self.normK+np.mat(np.diagflat(self.alpha*1./bres.f))
        bres.nabQ=np.linalg.inv(nabTemp)
        bres.diff=self.kernel * self.f - self.bd.data
        return bres       
    def minimizeA(self, alpha):
        self.updateAlpha(alpha)
        iter = 0
        while True:                
            self.mu = 1E-10
            while True:
                du = self.calcdu()           	
                if (self.calcStepSize(du) < self.convStep):
                    self.u += du		    
                    break
                else:
                    self.mu *= 10
            t = self.calcConv()
            if (t < self.convMin):
                print "Final t -", t                
                break
            if (iter > self.maxIter):
                print "Not Converged"
                break
        return self.sumResults()  
